# 使用数据库
import os
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.backends import default_backend
from cryptography.exceptions import InvalidSignature
from phe import paillier
import tenseal as ts
import json
import base64
import hashlib
from datetime import datetime, timedelta
import psycopg2
from psycopg2.extras import Json


class DataOwnershipSystem:
    def __init__(self, db_params=None):
        # 初始化数据库连接
        self.db_params = db_params or {
            'dbname': 'db427',
            'user': 'root',
            'password': 'root',
            'host': 'localhost',
            'port': '5432'
        }
        self.conn = psycopg2.connect(**self.db_params)

        # 初始化同态加密密钥
        self.paillier_public_key, self.paillier_private_key = paillier.generate_paillier_keypair(n_length=2048)

        # 初始化TenSEAL上下文
        context = ts.context(ts.SCHEME_TYPE.BFV, poly_modulus_degree=4096, plain_modulus=1032193)
        self.tenseal_context = context

        # 初始化数据库表
        self._init_db()

    def _init_db(self):
        """初始化数据库表结构"""
        with self.conn.cursor() as cur:
            # 创建所有权表
            cur.execute("""
                CREATE TABLE IF NOT EXISTS data_ownership (
                    data_hash TEXT PRIMARY KEY,
                    data TEXT NOT NULL,
                    public_key TEXT NOT NULL,
                    signature TEXT NOT NULL,
                    commitment TEXT NOT NULL,
                    blinding_factor TEXT NOT NULL,
                    registered_at TIMESTAMP NOT NULL,
                    metadata JSONB
                )
            """)

            # 创建使用权令牌表
            cur.execute("""
                CREATE TABLE IF NOT EXISTS usage_tokens (
                    token_id TEXT PRIMARY KEY,
                    encrypted_token JSONB NOT NULL,
                    data_hash TEXT NOT NULL,
                    seller_public_key TEXT NOT NULL,
                    buyer_public_key TEXT NOT NULL,
                    valid_until TIMESTAMP NOT NULL,
                    generated_at TIMESTAMP NOT NULL,
                    parent_token TEXT REFERENCES usage_tokens(token_id),
                    is_active BOOLEAN DEFAULT TRUE,
                    metadata JSONB
                )
            """)

            # 创建交易记录表
            cur.execute("""
                CREATE TABLE IF NOT EXISTS transactions (
                    transaction_id SERIAL PRIMARY KEY,
                    transaction_type TEXT NOT NULL,
                    transaction_data JSONB NOT NULL,
                    transaction_hash TEXT NOT NULL UNIQUE,
                    previous_hash TEXT,
                    timestamp TIMESTAMP NOT NULL,
                    block_height INTEGER
                )
            """)

            # 创建哈希链索引
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_transactions_hash ON transactions(transaction_hash)
            """)
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_transactions_prev_hash ON transactions(previous_hash)
            """)

            self.conn.commit()

    def generate_key_pair(self):
        """生成ECDSA密钥对用于所有权证明"""
        private_key = ec.generate_private_key(ec.SECP384R1(), default_backend())
        public_key = private_key.public_key()
        return private_key, public_key

    def create_pedersen_commitment(self, data, blinding_factor=None):
        """创建Pedersen承诺"""
        if blinding_factor is None:
            blinding_factor = os.urandom(32)

        # 简化的Pedersen承诺实现 (实际应使用椭圆曲线点)
        h = hashlib.sha256(blinding_factor).digest()
        data_hash = hashlib.sha256(data.encode()).digest()
        commitment = bytes(a ^ b for a, b in zip(h, data_hash))
        return commitment, blinding_factor

    def verify_ownership(self, data, signature, public_key):
        """验证数据所有权"""
        try:
            # 验证签名
            public_key.verify(
                base64.b64decode(signature),
                data.encode(),
                ec.ECDSA(hashes.SHA256()))

            # 验证所有权记录存在
            data_hash = hashlib.sha256(data.encode()).hexdigest()

            with self.conn.cursor() as cur:
                cur.execute(
                    "SELECT commitment, blinding_factor FROM data_ownership WHERE data_hash = %s",
                    (data_hash,)
                )
                result = cur.fetchone()
                if not result:
                    return False

                # 验证Pedersen承诺
                commitment, blinding_factor = result
                computed_commitment, _ = self.create_pedersen_commitment(
                    data,
                    base64.b64decode(blinding_factor))

                if base64.b64encode(computed_commitment).decode() != commitment:
                    return False

                return True
        except InvalidSignature:
            return False
        except Exception as e:
            print(f"Verification error: {e}")
            return False

    def _encrypt_with_paillier(self, data):
        """使用Paillier加密数据 - 确保输入为JSON字符串"""
        # 确保输入是JSON字符串
        if isinstance(data, dict):
            data_str = json.dumps(data, sort_keys=True)
        else:
            data_str = data

        # 将字符串分成小块(每块最多100字符)
        chunk_size = 100
        chunks = [data_str[i:i + chunk_size] for i in range(0, len(data_str), chunk_size)]

        encrypted_chunks = []
        for chunk in chunks:
            chunk_int = int.from_bytes(chunk.encode(), byteorder='big')
            encrypted = self.paillier_public_key.encrypt(chunk_int)
            encrypted_chunks.append({
                'ciphertext': str(encrypted.ciphertext()),
                'exponent': encrypted.exponent
            })

        return encrypted_chunks

    def _decrypt_with_paillier(self, encrypted_chunks):
        """使用Paillier解密数据 - 返回原始字典"""
        decrypted_chunks = []
        for chunk in encrypted_chunks:
            ciphertext = int(chunk['ciphertext'])
            exponent = chunk['exponent']
            encrypted = paillier.EncryptedNumber(
                self.paillier_public_key,
                ciphertext,
                exponent
            )
            decrypted_int = self.paillier_private_key.decrypt(encrypted)
            byte_length = (decrypted_int.bit_length() + 7) // 8
            decrypted_bytes = decrypted_int.to_bytes(byte_length, byteorder='big')
            decrypted_chunks.append(decrypted_bytes.decode())

        full_data = ''.join(decrypted_chunks)
        return json.loads(full_data)  # 返回字典而非字符串



    def register_ownership(self, data, private_key):
        """注册数据所有权 - 使用PostgreSQL存储"""
        #这里给出的所有权注册仅为机器自动生成（主要为了验证交易链完整性），手动注册数据所有权前面已经展示
        try:
            # 1. 创建数据签名
            signature = private_key.sign(data.encode(), ec.ECDSA(hashes.SHA256()))

            # 2. 创建Pedersen承诺
            commitment, blinding_factor = self.create_pedersen_commitment(data)

            # 3. 准备所有权记录
            public_key = private_key.public_key()
            public_key_bytes = public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )

            data_hash = hashlib.sha256(data.encode()).hexdigest()
            registered_at = datetime.utcnow()

            # 4. 存储到PostgreSQL
            with self.conn.cursor() as cur:
                # 检查所有权唯一性，哈希算法作用：通过 data_hash 字段唯一标识数据，确保相同数据的哈希值唯一。通常采用 SHA-256 等抗碰撞哈希算法，保证不同数据生成相同哈希值的概率极低
                #防重复注册：通过哈希值作为唯一索引，避免同一数据被多次注册，确保数据所有权的唯一性
                cur.execute(
                    "SELECT 1 FROM data_ownership WHERE data_hash = %s",
                    (data_hash,)
                )
                if cur.fetchone():
                    raise ValueError("Data ownership already registered")

                # 插入所有权记录
                cur.execute("""
                    INSERT INTO data_ownership (
                        data_hash, data, public_key, signature,
                        commitment, blinding_factor, registered_at
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                """, (
                    data_hash,
                    data,
                    public_key_bytes.decode(),#public_key 是非对称加密中的公钥，用于验证签名。通常由用户生成（如 RSA 或 ECC 密钥对），公钥公开存储
                    base64.b64encode(signature).decode(),#签名生成：数据所有者使用私钥对 data_hash 或其他关键信息生成签名，确保数据来源真实性和不可抵赖性，验证时需用公钥解密签名，与原始数据哈希比对，确认数据未被篡改且签名合法
                    base64.b64encode(commitment).decode(),#承诺值
                    base64.b64encode(blinding_factor).decode(),#盲化因子 #使用 blinding_factor（盲化因子）对原始数据计算生成 commitment（承诺值），如 Pedersen 承诺：commitment = H(data || blinding_factor)；揭示阶段：事后可通过公开 data 和 blinding_factor 验证承诺值是否匹配，但承诺阶段不泄露原始数据
                    registered_at
                ))

                # 添加到交易记录 交易链与防篡改机制
                transaction_data = {
                    'data_hash': data_hash,
                    'public_key': public_key_bytes.decode(),
                    'timestamp': registered_at.isoformat()
                }
                #链式结构：每笔交易记录包含 data_hash 和 timestamp，通过哈希链链接，确保历史记录不可篡改
                self._add_to_transaction_chain('REGISTER', transaction_data)

                self.conn.commit()

            return {
                'data_hash': data_hash,
                'public_key': public_key_bytes.decode(),
                'timestamp': registered_at.isoformat() #时间戳作用：记录操作时间，防止重放攻击
            }
        except Exception as e:
            self.conn.rollback()
            raise e


    def generate_usage_token(self, data, seller_private_key, buyer_public_key, validity_period=365):
        """生成使用权令牌 - 使用PostgreSQL存储"""
        try:
            # 1. 验证卖方所有权
            data_hash = hashlib.sha256(data.encode()).hexdigest()

            with self.conn.cursor() as cur:
                cur.execute(
                    "SELECT public_key, signature FROM data_ownership WHERE data_hash = %s",
                    (data_hash,)
                )
                result = cur.fetchone()
                if not result:
                    raise ValueError("Data ownership not registered")

                public_key_pem, signature = result
                public_key = serialization.load_pem_public_key(
                    public_key_pem.encode(),
                    backend=default_backend()
                )

                if not self.verify_ownership(data, signature, public_key):
                    raise ValueError("Seller ownership verification failed")

            # 2. 创建令牌元数据
            token_metadata = {
                'data_hash': data_hash,
                'seller_public_key': seller_private_key.public_key().public_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo
                ).decode(),
                'buyer_public_key': buyer_public_key.public_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo
                ).decode(),
                'valid_until': (datetime.utcnow() + timedelta(days=validity_period)).isoformat(),
                'generation_time': datetime.utcnow().isoformat(),
                'parent_token': None
            }

            # 3. 加密令牌
            encrypted_token = self._encrypt_with_paillier(token_metadata)

            # 4. 存储到PostgreSQL
            token_id = hashlib.sha256(os.urandom(32)).hexdigest()
            '''
            encrypted_token，字段存储对称加密后的令牌数据（如AES-GCM加密），需在应用层完成加密后通过Json()序列化存储。典型实现使用cryptography库的Fernet模块（基于AES-128/CBC+HMAC），同时提供机密性和完整性验证
            data_hash,存储原始数据的哈希值（可能为SHA-256或HMAC），用于验证数据在传输/存储过程中是否被篡改
            seller_public_key和buyer_public_key涉及非对称加密体系（如RSA/ECC），身份验证：通过数字签名验证双方交易身份；密钥交换：后续流程中会生成会话密钥（如ECDH算法）
            #加密字段存储原理（上面部分）
            '''
            with self.conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO usage_tokens (
                        token_id, encrypted_token, data_hash,
                        seller_public_key, buyer_public_key,
                        valid_until, generated_at, metadata
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    token_id,
                    Json(encrypted_token),
                    data_hash,
                    token_metadata['seller_public_key'],
                    token_metadata['buyer_public_key'],
                    datetime.fromisoformat(token_metadata['valid_until']),
                    datetime.fromisoformat(token_metadata['generation_time']),
                    Json(token_metadata)
                ))

                '''
                 添加到交易记录（交易链完整性原理体现）
                 哈希链结构：每次交易记录包含前一笔的哈希值，形成不可篡改的链式结构（类似区块链），修改任一记录会导致后续所有哈希值失效
                 防篡改机制：data_hash确保交易内容完整性，结合时间戳字段（generated_at）可防御重放攻击
                 '''
                self._add_to_transaction_chain('TOKEN_GENERATE', {
                    'token_id': token_id,
                    'data_hash': data_hash,
                    'seller': token_metadata['seller_public_key'],
                    'buyer': token_metadata['buyer_public_key']
                })

                self.conn.commit()

            return token_id
        except Exception as e:
            self.conn.rollback()
            raise e

    def resell_data(self, token_id, reseller_private_key, new_buyer_public_key, validity_period=30):
        """转售数据 - 使用PostgreSQL存储"""
        try:
            with self.conn.cursor() as cur:
                # 1. 获取原令牌
                cur.execute("""
                    SELECT encrypted_token, metadata, valid_until, buyer_public_key
                    FROM usage_tokens
                    WHERE token_id = %s AND is_active = TRUE
                """, (token_id,))

                result = cur.fetchone()
                if not result:
                    raise ValueError("Invalid token ID")

                encrypted_token, metadata, valid_until, buyer_public_key = result

                # 2. 验证转售者
                reseller_public_key = reseller_private_key.public_key().public_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo
                ).decode()

                if reseller_public_key != buyer_public_key:
                    raise ValueError("Reseller is not the current owner")

                # 3. 检查令牌是否过期
                if valid_until < datetime.utcnow():
                    raise ValueError("Token has expired")

                # 4. 创建新令牌(权限衰减)
                new_validity_period = min(
                    validity_period,
                    (valid_until - datetime.utcnow()).days
                )

                new_token_metadata = {
                    'data_hash': metadata['data_hash'],
                    'seller_public_key': reseller_public_key,
                    'buyer_public_key': new_buyer_public_key.public_bytes(
                        encoding=serialization.Encoding.PEM,
                        format=serialization.PublicFormat.SubjectPublicKeyInfo
                    ).decode(),
                    'valid_until': (datetime.utcnow() + timedelta(days=new_validity_period)).isoformat(),
                    'generation_time': datetime.utcnow().isoformat(),
                    'parent_token': token_id
                }

                # 5. 加密新令牌  同态加密
                encrypted_token = self._encrypt_with_paillier(new_token_metadata)

                # 6. 存储新令牌  哈希链式存储
                new_token_id = hashlib.sha256(os.urandom(32)).hexdigest()
                '''
                密钥管理：
                buyer_public_key和seller_public_key的存储表明使用了非对称加密体系，公钥可公开分发，私钥需保密

                '''
                cur.execute("""
                    INSERT INTO usage_tokens (
                        token_id, encrypted_token, data_hash,
                        seller_public_key, buyer_public_key,
                        valid_until, generated_at, parent_token, metadata
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    new_token_id,
                    Json(encrypted_token),
                    new_token_metadata['data_hash'],
                    new_token_metadata['seller_public_key'],
                    new_token_metadata['buyer_public_key'],
                    datetime.fromisoformat(new_token_metadata['valid_until']),
                    datetime.fromisoformat(new_token_metadata['generation_time']),
                    new_token_metadata['parent_token'],
                    Json(new_token_metadata)
                ))

                # 7. 添加到交易记录
                '''
                主要原理与前面相同，交易记录以链式结构存储，每个交易包含父令牌ID（parent_token），通过哈希值链接前后交易，确保交易历史的不可篡改性
                身份认证：reseller和new_buyer的公钥用于验证交易参与者身份，防止伪造交易
                元数据存储Json(new_token_metadata)，元数据可能包含敏感信息（如密钥参数），JSON序列化后需加密存储或结合数字签名，防止数据泄露或篡改，密钥分离原则：元数据中的公钥与私钥分离存储，遵循非对称加密的最小权限原则
                '''
                self._add_to_transaction_chain('TOKEN_RESALE', {
                    'parent_token_id': token_id,
                    'new_token_id': new_token_id,
                    'reseller': reseller_public_key,
                    'new_buyer': new_token_metadata['buyer_public_key']
                })

                self.conn.commit()

            return new_token_id
        except Exception as e:
            self.conn.rollback()
            raise e

    def _add_to_transaction_chain(self, transaction_type, transaction_data):
        """添加交易记录到防篡改链 - 修正哈希计算"""
        try:
            with self.conn.cursor() as cur:
                # 哈希链构建，获取最近交易的哈希值作为当前交易的previous_hash字段，形成不可逆的链式结构，任何交易篡改都会破坏后续哈希值（雪崩效应）
                cur.execute("""
                    SELECT transaction_hash FROM transactions
                    ORDER BY transaction_id DESC LIMIT 1
                """)
                last_tx = cur.fetchone()
                previous_hash = last_tx[0] if last_tx else None

                # 获取当前时间戳（统一使用一个时间点）
                current_time = datetime.utcnow()
                timestamp_str = current_time.isoformat() + 'Z'

                # 构建标准化的交易记录（固定字段顺序）
                transaction = {
                    'type': transaction_type,
                    'data': transaction_data,
                    'timestamp': timestamp_str,
                    'previous_hash': previous_hash
                }

                # 计算哈希（确保字段顺序一致）
                transaction_str = json.dumps(transaction, sort_keys=True, separators=(',', ':'))
                current_hash = hashlib.sha256(transaction_str.encode()).hexdigest()

                # 数据库的存储设计
                cur.execute("""
                    INSERT INTO transactions (
                        transaction_type, transaction_data,
                        transaction_hash, previous_hash, timestamp
                    ) VALUES (%s, %s, %s, %s, %s)
                    RETURNING transaction_id
                """, (
                    transaction_type,
                    Json(transaction),  # 存储完整结构，全量数据存储，通过Json(transaction)存储完整交易结构，任何字段修改都会导致存储的哈希值与重新计算结果不一致
                    current_hash,
                    previous_hash,#同时保存current_hash和previous_hash，形成可验证的哈希链
                    current_time  # 使用相同的时间戳，也要注意数据库时间精度问题，避免影响后续验证
                ))

                tx_id = cur.fetchone()[0]
                self.conn.commit()
                #原子性提交，要么写入完整所有字段，要么完全回滚，防止出现“半完成交易”破坏哈希链连续性
                print(f"已添加交易#{tx_id} 类型:{transaction_type}")
                print(f"计算哈希的内容: {transaction_str}")
                print(f"存储的哈希: {current_hash}")
                #明文展示哈希计算内容，便于事后审计，符合要求中密码学系统的可公开验证性
                return current_hash
        except Exception as e:
            self.conn.rollback()
            print(f"添加交易失败: {str(e)}")
            raise



    def verify_transaction_chain(self):
        """验证交易链完整性 - 20250421修正版"""
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    SELECT transaction_id, transaction_hash, previous_hash,
                           transaction_data, timestamp
                    FROM transactions
                    ORDER BY transaction_id
                """)
                transactions = cur.fetchall()

                print(f"\n验证交易链，共{len(transactions)}笔交易")

                if not transactions:
                    print("空交易链视为有效")
                    return True

                # 验证第一笔交易
                tx_id, tx_hash, prev_hash, data, ts = transactions[0]
                print(f"\n验证交易#{tx_id} 类型:{data['type']}")

                if prev_hash is not None:
                    print("错误: 第一笔交易的前一个哈希不为None")
                    return False

                # 构建期望的哈希内容
                expected_data = {
                    'type': data['type'],
                    'data': data['data'],
                    'timestamp': ts.isoformat() + 'Z',  # 保持与存储时相同格式
                    'previous_hash': None
                }
                tx_str = json.dumps(expected_data, sort_keys=True, separators=(',', ':'))
                computed_hash = hashlib.sha256(tx_str.encode()).hexdigest()

                print(f"计算哈希的内容: {tx_str}")
                print(f"存储的哈希: {tx_hash}")
                print(f"计算的哈希: {computed_hash}")

                if computed_hash != tx_hash:
                    print(f"哈希不匹配! 差异:")
                    print(f"存储的内容: {data}")
                    print(f"计算的内容: {expected_data}")
                    return False

                # 验证后续交易...
                # ...（保持原有逻辑）

                return True
        except Exception as e:
            print(f"验证交易链时出错: {str(e)}")
            return False


    def close(self):
        """关闭数据库连接"""
        # 如果self.conn存在，则关闭数据库连接
        if self.conn:
            self.conn.close()